'''
author:        wangchenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2023-12-23
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

'''
    将中文自动补全插件整理成词典
'''
from xpinyin import Pinyin as 拼音库类
import csv
import pickle

声母表 = ['b', 'p', 'm', 'f', 'd', 't', 'n', 'l', 'g', 'k', 'h', 
       'j', 'q', 'x', 'zh', 'ch', 'sh','r', 'z', 'c', 's', 'y', 'w']
韵母表 = ['a', 'o', 'e', 'i', 'u', 'v', 'ai', 'ei', 'ui', 'ao', 'ou',
       'iu', 'ie', 've', 'er', 'an', 'en', 'in', 'un', 'vn', 
       'ang', 'eng', 'ing', 'ong', 'ia', 'ian', 'iang', 'iao', 'iong',
       'ua', 'uai', 'uan', 'uang', 'uo', 'van']

class 字类:
    汉字:str
    声母:str
    韵母:str
    声调:int
    字频:int
    def __init__(self, 汉字, 声母, 韵母, 声调, 字频) -> None:
        self.汉字 = 汉字
        self.声母 = 声母
        self.韵母 = 韵母
        self.声调 = 声调
        self.字频 = 字频
    
    def 转列表(self):
        return [self.汉字, self.声母, self.韵母, self.声调, self.字频]
    
    def 转标准拼音(self): 
        # 将当前字的拼音转为标准拼音
        韵母字 = self.韵母
        if(self.声母 in ['j', 'q', 'x', 'y']): 
            if(self.韵母 == 'v'): 
                韵母字 = 'u'
            elif(self.韵母 == 've'): 
                韵母字 = 'ue'
            elif(self.韵母 == 'vn'): 
                韵母字 = 'un'
            elif(self.韵母 == 'van'): 
                韵母字 = 'uan'
        return self.声母 + 韵母字

def 转标准拼音(字:字类) -> str: 
    return 字.转标准拼音()

def 转小鹤(字:字类) -> str: 
    声母 = 字.声母
    if(声母 == 'zh'): 
        声母 = 'v'
    elif(声母 == 'ch'): 
        声母 = 'i'
    elif(声母 == 'sh'): 
        声母 = 'u'
    
    韵母 = 字.韵母
    转换表 = {
        'ai':'d',
        'ei':'w',
        'ui':'v',
        'ao':'c',
        'ou':'z',
        'iu':'q',
        'ie':'p',
        've':'t',
        'an':'j',
        'en':'f',
        'in':'b',
        'un':'y',
        'vn':'y',
        'ang':'h',
        'eng':'g',
        'ing':'k',
        'ong':'s',
        'ia':'x',
        'ian':'m',
        'iang':'l',
        'iao':'n',
        'iong':'s',
        'ua':'x',
        'uai':'k',
        'uan':'r',
        'uang':'l',
        'uo':'o',
        'van':'r'
    }
    if(声母 == ''): 
        if(len(韵母) == 1): 
            return 韵母 * 2
        elif(len(韵母) > 2): 
            return 韵母[0] + 转换表[韵母]
        else: 
            return 韵母
    if(len(韵母) > 1): 
        韵母 = 转换表[韵母]
    return 声母 + 韵母
    

class 词类: 
    字列:list
    词频:int
    def __init__(self, 字列, 词频) -> None:
        self.字列 = 字列
        self.词频 = 词频

def 拼音串拆分(拼音串:str): 
    # 分离声调
    try:
        声调 = int(拼音串[-1])
    except: 
        print(拼音串)
        assert(0)
    拼音串 = 拼音串[:-1]

    # 分离声母
    if(拼音串[0] in 声母表):
        # 有声母
        if(拼音串[:2] in ['zh', 'ch', 'sh']): 
            # 声母是zh, ch, sh
            声母 = 拼音串[:2]
            韵母 = 拼音串[2:]
        else: 
            声母 = 拼音串[0]
            韵母 = 拼音串[1:]
    else:
        # 无声母
        声母 = ''
        韵母 = 拼音串
    
    # 韵母标准化
    if(韵母 == 'u'): 
        if(声母 in ['j', 'q', 'x', 'y']): 
            韵母 = 'v'
    if(韵母 == 'ue'): 
        if(声母 in ['j', 'q', 'x', 'y', 'n']): 
            韵母 = 've'
    if(韵母 == 'un'): 
        if(声母 in ['j', 'q', 'x', 'y']): 
            韵母 = 'vn'
    if(韵母 == 'uan'): 
        if(声母 in ['j', 'q', 'x', 'y']): 
            韵母 = 'van'
    
    return (声母, 韵母, 声调)
    

def xpinyin转字(汉字:str, 字频:int):
    拼音库 = 拼音库类()
    所有拼音列表 = 拼音库.get_pinyins(汉字, tone_marks='numbers')

    字表 = []
    for 当前拼音 in 所有拼音列表: 
        while(not 当前拼音[-1] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']): 
            当前拼音 = 当前拼音[:-1]
            if(len(当前拼音) == 0): 
                break
        if(len(当前拼音) == 0): 
            continue
        声母, 韵母, 声调 = 拼音串拆分(当前拼音)
        字表.append(字类(汉字, 声母, 韵母, 声调, 字频))
    
    return 字表

def 本地字典转字(汉字:str, 预选拼音:str, 字频:int): 
    # 将当前汉字通过xpinyin转成字
    字表 = []
    # 1. 从汉字得到字表 
    预选字表 = xpinyin转字(汉字, 字频)

    # 2. 按照预选拼音选出合适的字
    for 字 in 预选字表: 
        当前拼音 = 转标准拼音(字)
        if(预选拼音 == 当前拼音): 
            字表.append(字)
    
    if(len(预选字表) == 0): 
        声母, 韵母, 声调 = 拼音串拆分(预选拼音 + '0')
        字表 = [字类(汉字, 声母, 韵母, 声调, 字频)]
    return 字表

def 本地词典转词(词组:str, 拼音列表:list[str], 词频): 
    字列 = []
    for j in range(len(词组)): 
        预选字表 = xpinyin转字(词组[j])
        当前拼音 = 拼音列表[j]
        有字选取 = False
        for 字 in 预选字表: 
            if(转标准拼音(字) == 当前拼音): 
                字列.append(字)
                有字选取 = True
                break
        if(not 有字选取): 
            raise ValueError("词组" + 词组 + "没有合适的拼音")
    return 词类(字列, 词频)


def 程序_测试xpinyin转字():
    当前字 =  xpinyin转字("落", 0)
    for 字 in 当前字: 
        print(字.转标准拼音())
        print(字.声调)

def 程序_整理字库(): 
    字库 = set()
    with open("字库.txt", 'r') as fp: 
        reader = csv.reader(fp, delimiter='\t')
        for lines in reader: 
            lines_striped = [s.strip() for s in lines]
            字表 = 本地字典转字(lines_striped[0], lines_striped[1], int(lines_striped[2]))
            for 字 in 字表: 
                字库.add(字)
                print(字.汉字, 字.转标准拼音(), 字.声调)
    with open("字库.pkl", 'wb') as fp:
        pickle.dump(字库, fp)

def 程序_生癖字(): 
    with open("字库.pkl", "rb") as fp: 
        字库 = pickle.load(fp)
    
    for 字 in 字库: 
        if(字.声调 == 0): 
            print(字.转列表())

def 程序_生成小鹤词典(): 
    with open("字库.pkl", "rb") as fp: 
        字库 = pickle.load(fp)

    字库列表 = list(字库)
    字库列表.sort(key=lambda 字: (字.声母, 字.韵母, -字.字频))

    with open("xiaohe-wcy.dict.txt", 'w') as fp:
        for 字 in 字库列表: 
            拼音 = 转小鹤(字)
            首字母 = 拼音[0].upper()
            拼音 = 首字母 + 拼音[1:]
            print("%s\t %s\t %d"%(字.汉字, 拼音, 字.字频), end='\r')
            print("%s\t %s\t %d"%(字.汉字, 拼音, 字.字频), file=fp)

def 程序_生成声调小鹤词典(): 
    with open("字库.pkl", "rb") as fp: 
        字库 = pickle.load(fp)

    字库列表 = list(字库)
    字库列表.sort(key=lambda 字: (字.声母, 字.韵母, -字.字频))

    声调键 = [' ', 'j', 'h', 'g', 'f','i']
    with open("tune-xiaohe-wcy.dict.txt", 'w') as fp:
        for 字 in 字库列表: 
            拼音 = 转小鹤(字)
            首字母 = 拼音[0].upper()
            拼音 = 首字母 + 拼音[1:] + 声调键[字.声调]
            print("%s\t %s\t %d"%(字.汉字, 拼音, 字.字频), end='\r')
            print("%s\t %s\t %d"%(字.汉字, 拼音, 字.字频), file=fp)




if __name__ == "__main__": 
    # 程序_测试xpinyin转字()
    # 程序_整理字库()
    # 程序_生癖字()
    # 程序_生成小鹤词典()
    程序_生成声调小鹤词典()